#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
WSL_Manager
新增：最小化到托盘 / 双击托盘恢复 / 自定义 ICO
"""
import sys, pathlib, subprocess, datetime, shutil
from PySide6.QtCore import (Qt, QThread, QObject, Signal, Slot,
                            QMetaObject, Q_ARG, QTimer, QProcess, QEvent)
from PySide6.QtWidgets import *
from PySide6.QtGui import QIcon
from PySide6.QtWidgets import QStyle
startupinfo = subprocess.STARTUPINFO()
startupinfo.dwFlags |= subprocess.STARTF_USESHOWWINDOW
startupinfo.wShowWindow = subprocess.SW_HIDE
# ---------------- 后端 ----------------
class WslBackend(QObject):
    refresh = Signal(list)
    log = Signal(str)
    error = Signal(str)
    progress = Signal(int)
    progress_msg = Signal(str)
    cmd_out = Signal(str)

    def _run(self, cmd: list):
        try:
            out = subprocess.check_output(cmd, stderr=subprocess.STDOUT,startupinfo=startupinfo)
            return True, out.decode("utf-16-le",errors="ignore")
        except subprocess.CalledProcessError as e:
            return False, e.output.decode("utf-16-le",errors="ignore")
        except FileNotFoundError:
            return False, "找不到 wsl.exe，请确认 Windows 子系统已安装。"

    @Slot()
    def list_distros(self):
        ok, txt = self._run(["wsl.exe", "-l", "-v"])
        if not ok:
            self.error.emit(txt); return
        data = []
        for ln in txt.splitlines()[1:]:
            ln = ln.lstrip("*").strip()
            parts = [p for p in ln.split() if p]
            if len(parts) >= 3:
                data.append((parts[0], parts[-2], parts[-1]))
        self.refresh.emit(data)

    @Slot(str)
    def set_default(self, distro: str):
        ok, txt = self._run(["wsl.exe", "--set-default", distro])
        self.log.emit(f"已设置 {distro} 为默认发行版" if ok else "设置默认失败:\n" + txt)

    @Slot(str, str)
    def resize_memory(self, distro: str, mem: str):
        ok, txt = self._run(["wsl.exe", "--manage", distro, "--resize", mem])
        self.log.emit(f"已调整 {distro} 内存为 {mem}" if ok else "调整内存失败:\n" + txt)

    @Slot(str, str)
    def move_distro(self, distro: str, new_path: str):
        ok, txt = self._run(["wsl.exe", "--manage", distro, "--move", new_path])
        self.log.emit(f"已移动 {distro} 到 {new_path}" if ok else "移动失败:\n" + txt)

    @Slot(str, str)
    def set_default_user(self, distro: str, user: str):
        ok, txt = self._run(["wsl.exe", "--manage", distro, "--set-default-user", user])
        self.log.emit(f"已设置 {distro} 默认用户为 {user}" if ok else "设置用户失败:\n" + txt)

    @Slot(str, str)
    def run_as_user(self, distro: str, user: str):
        subprocess.Popen(["cmd", "/c", "start", "wsl", "-d", distro, "--user", user])
        self.log.emit(f"已启动 {distro} 用户={user}")

    @Slot(str)
    def terminate(self, distro: str):
        ok, txt = self._run(["wsl.exe", "--terminate", distro])
        self.log.emit(f"{distro} 已终止" if ok else txt)

    @Slot(str, str)
    def export_distro(self, distro: str, path: str):
        self._do_progress_cmd(["wsl.exe", "--export", distro, path], f"正在导出  {distro} …")

    @Slot(str, str, str)
    def import_distro(self, distro: str, install_dir: str, tar_path: str):
        self._do_progress_cmd(["wsl.exe", "--import", distro, install_dir, tar_path], f"正在导入  {distro} …")

    @Slot(str)
    def unregister(self, distro: str):
        ok, txt = self._run(["wsl.exe", "--unregister", distro])
        self.log.emit(f"{distro} 已卸载" if ok else txt)

    @Slot(str, str)
    def rename_distro(self, old_name: str, new_name: str):
        tmp_tar = pathlib.Path.home() / "WSLTemp" / f"{old_name}_{datetime.datetime.now():%Y%m%d%H%M%S}.tar"
        tmp_tar.parent.mkdir(exist_ok=True)
        self.progress_msg.emit(f"重命名【{old_name}→{new_name}】第一步：导出中…")
        self.progress.emit(0)
        ok, txt = self._run(["wsl.exe", "--export", old_name, str(tmp_tar)])
        if not ok:
            self.error.emit("导出失败:\n" + txt); return
        self.progress_msg.emit("第二步：注销旧发行版…")
        ok, txt = self._run(["wsl.exe", "--unregister", old_name])
        if not ok:
            self.error.emit("注销失败:\n" + txt); return
        install_dir = pathlib.Path.home() / "WSLImport" / new_name
        install_dir.mkdir(parents=True, exist_ok=True)
        self.progress_msg.emit("第三步：导入中…")
        self.progress.emit(0)
        ok, txt = self._run(["wsl.exe", "--import", new_name, str(install_dir), str(tmp_tar)])
        if not ok:
            self.error.emit("导入失败:\n" + txt); return
        tmp_tar.unlink(missing_ok=True)
        self.log.emit(f"重命名完成：{old_name} → {new_name}")
        self.progress.emit(100)

    def _do_progress_cmd(self, cmd: list, msg: str):
        self.progress_msg.emit(msg); self.progress.emit(0)
        proc = QProcess()
        proc.setProgram(cmd[0]); proc.setArguments(cmd[1:])
        proc.setProcessChannelMode(QProcess.SeparateChannels)

        def on_stderr():
            for tok in bytes(proc.readAllStandardError()).decode("utf-16-le",errors="ignore").split():
                if tok.endswith("%"):
                    try:
                        self.progress.emit(int(tok[:-1]))
                    except ValueError:
                        pass
        proc.readyReadStandardError.connect(on_stderr)
        proc.start()
        proc.waitForFinished(-1)
        ok = proc.exitCode() == 0
        self.progress.emit(100 if ok else -1)
        self.log.emit(msg.replace("正在", "完成")) if ok else self.error.emit(bytes(proc.readAllStandardError()).decode("utf-16-le",errors="ignore"))

    @Slot(str)
    def run_cmd(self, raw: str):
        if not raw.strip():
            self.error.emit("命令不能为空"); return
        self.progress_msg.emit(f"执行: {raw}"); self.progress.emit(0)
        proc = QProcess()
        proc.setProgram("cmd"); proc.setArguments(["/C", raw])
        proc.setProcessChannelMode(QProcess.SeparateChannels)
        proc.readyReadStandardOutput.connect(lambda: self.cmd_out.emit(bytes(proc.readAllStandardOutput()).decode("utf-16-le",errors="ignore")))
        proc.readyReadStandardError.connect(lambda: self.cmd_out.emit(bytes(proc.readAllStandardError()).decode("utf-16-le",errors="ignore")))
        proc.start()
        proc.waitForFinished(-1)
        self.progress.emit(100 if proc.exitCode() == 0 else -1)
        self.log.emit("命令执行完成" if proc.exitCode() == 0 else "命令执行失败")


# ---------------- 主窗 ----------------
class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("WSL_Manager")
        self.resize(720, 620)
        self.setWindowFlags(Qt.WindowCloseButtonHint | Qt.MSWindowsFixedSizeDialogHint)

        # -------------- 托盘图标 --------------
        self._init_tray()
        central = QWidget(); self.setCentralWidget(central)
        lay = QVBoxLayout(central)
        # 表格
        self.table = QTableWidget(0, 3)
        self.table.setHorizontalHeaderLabels(["发行版", "状态", "WSL 版本"])
        self.table.horizontalHeader().setStretchLastSection(True)
        self.table.setSelectionBehavior(QTableWidget.SelectRows)
        self.table.setEditTriggers(QTableWidget.NoEditTriggers)
        self.table.doubleClicked.connect(self.enter)
        self.table.setContextMenuPolicy(Qt.CustomContextMenu)
        self.table.customContextMenuRequested.connect(self.open_menu)
        lay.addWidget(self.table)

        # 按钮栏
        btn_bar = QHBoxLayout()
        self._buttons = []
        for txt, slot in (("刷新", self.refresh), ("启动", self.enter), ("终止", self.terminate),
                          ("备份", self.backup), ("导入", self.import_distro), ("关机", self.shutdown),
                          ("卸载", self.unregister)):
            btn = QPushButton(txt); btn.clicked.connect(slot)
            btn_bar.addWidget(btn); self._buttons.append(btn)
        lay.addLayout(btn_bar)

        # 用户运行栏
        user_bar = QHBoxLayout()
        user_bar.addWidget(QLabel("用户:"))
        self.user_line = QLineEdit(); self.user_line.setPlaceholderText("root 或普通用户名")
        self.user_btn = QPushButton("用户运行"); self.user_btn.clicked.connect(self.run_as_user)
        user_bar.addWidget(self.user_line); user_bar.addWidget(self.user_btn)
        lay.addLayout(user_bar)

        # 自定义命令
        cmd_bar = QHBoxLayout()
        cmd_bar.addWidget(QLabel("命令:"))
        self.cmd_line = QLineEdit(); self.cmd_line.setPlaceholderText('例：wsl -d Ubuntu -- apt update')
        self.run_btn = QPushButton("执行"); self.run_btn.clicked.connect(self.run_custom_cmd)
        cmd_bar.addWidget(self.cmd_line); cmd_bar.addWidget(self.run_btn)
        lay.addLayout(cmd_bar)

        # 日志
        self.log = QPlainTextEdit(); self.log.setMaximumHeight(100); self.log.setReadOnly(True)
        lay.addWidget(self.log)

        # 状态栏 + 进度条
        self.status = QLabel("就绪")
        self.progress_bar = QProgressBar(); self.progress_bar.setVisible(False); self.progress_bar.setFixedWidth(140)
        self.statusBar().addWidget(self.status, 1); self.statusBar().addPermanentWidget(self.progress_bar)

        # 线程
        self.thread = QThread()
        self.backend = WslBackend(); self.backend.moveToThread(self.thread)
        self.backend.refresh.connect(self.fill_table)
        self.backend.log.connect(self.append_log)
        self.backend.error.connect(lambda e: QMessageBox.critical(self, "错误", e))
        self.backend.progress.connect(self.on_progress)
        self.backend.progress_msg.connect(self.on_progress_msg)
        self.backend.cmd_out.connect(self.append_log)
        self.thread.start()

        self.refresh()

    # -------------- 托盘相关 --------------
    def _init_tray(self):
        icon_file = pathlib.Path(__file__).with_name("Heart.ico")
        self.app_icon = QIcon(str(icon_file)) if icon_file.exists() else self.style().standardIcon(QStyle.SP_ComputerIcon)
        self.setWindowIcon(self.app_icon)
        self.tray = QSystemTrayIcon(self.app_icon, self)
        self.tray.setToolTip("WSL_Manager")
        menu = QMenu()
        show_action = menu.addAction("显示主界面")
        quit_action = menu.addAction("退出程序")
        show_action.triggered.connect(self._show_normal)
        quit_action.triggered.connect(qApp.quit)
        self.tray.setContextMenu(menu)
        self.tray.activated.connect(self._tray_activated)
        self.tray.show()

    def _tray_activated(self, reason):
        if reason == QSystemTrayIcon.DoubleClick:
            self._show_normal()

    def _show_normal(self):
        self.show(); self.raise_(); self.activateWindow()
        self.setWindowState(self.windowState() & ~Qt.WindowMinimized)

    def closeEvent(self, event):
        event.ignore(); self.hide()
        self.tray.showMessage("WSL_Manager", "程序已最小化到托盘，双击图标即可恢复。", QSystemTrayIcon.Information, 1500)

    def changeEvent(self, event):
        if event.type() == QEvent.WindowStateChange and self.isMinimized():
            self.hide(); event.accept()

    # -------------- 通用 --------------
    def append_log(self, txt):
        t = datetime.datetime.now().strftime("%H:%M:%S")
        self.log.appendPlainText(f"[{t}]  {txt}")

    def on_progress_msg(self, txt):
        self.statusBar().showMessage(txt); self.progress_bar.setVisible(True)

    def on_progress(self, pct):
        if pct < 0:
            self.progress_bar.setVisible(False); self.statusBar().showMessage("操作失败", 3000); return
        self.progress_bar.setValue(pct)
        if pct == 100:
            self.progress_bar.setVisible(False); self.statusBar().showMessage("就绪", 2000)

    def set_buttons(self, enable: bool):
        for btn in self._buttons:
            btn.setEnabled(enable)

    def refresh(self):
        self.set_buttons(False); self.statusBar().showMessage("刷新中…")
        QMetaObject.invokeMethod(self.backend, "list_distros", Qt.QueuedConnection)

    def fill_table(self, data):
        self.table.setRowCount(0)
        for name, state, ver in data:
            row = self.table.rowCount()
            self.table.insertRow(row)
            self.table.setItem(row, 0, QTableWidgetItem(name))
            self.table.setItem(row, 1, QTableWidgetItem(state))
            self.table.setItem(row, 2, QTableWidgetItem(ver))
        self.statusBar().showMessage(f"共 {len(data)} 个发行版")
        self.set_buttons(True)

    def current_distro(self):
        row = self.table.currentRow()
        if row < 0:
            QMessageBox.warning(self, "提示", "请先选中一个发行版！")
            return None
        return self.table.item(row, 0).text().strip()

    # -------------- 右键菜单 --------------
    def open_menu(self, pos):
        item = self.table.itemAt(pos)
        if not item:
            return
        distro = self.table.item(item.row(), 0).text()
        menu = QMenu()
        menu.addAction("重命名", self.rename_distro)
        menu.addAction("设为默认", self.set_default_distro)
        menu.addAction("以指定用户运行", self.run_selected_as_user)
        manage_menu = menu.addMenu("高级管理")
        manage_menu.addAction("调整内存", lambda: self.manage_resize(distro))
        manage_menu.addAction("移动目录", lambda: self.manage_move(distro))
        manage_menu.addAction("设置默认用户", lambda: self.manage_user(distro))
        menu.exec(self.table.mapToGlobal(pos))

    def set_default_distro(self):
        name = self.current_distro()
        if name:
            QMetaObject.invokeMethod(self.backend, "set_default", Qt.QueuedConnection, Q_ARG(str, name))

    def run_selected_as_user(self):
        distro = self.current_distro()
        if not distro:
            return
        user, ok = QInputDialog.getText(self, "以指定用户运行", "请输入用户名（如 root）：")
        if ok and user:
            QMetaObject.invokeMethod(self.backend, "run_as_user", Qt.QueuedConnection, Q_ARG(str, distro), Q_ARG(str, user))

    def rename_distro(self):
        old_name = self.current_distro()
        if not old_name:
            return
        new_name, ok = QInputDialog.getText(self, "重命名发行版", f"请输入新名称（原：{old_name}）：")
        if not ok or not new_name or new_name == old_name:
            return
        QMetaObject.invokeMethod(self.backend, "rename_distro", Qt.QueuedConnection, Q_ARG(str, old_name), Q_ARG(str, new_name))
        QTimer.singleShot(1500, self.refresh)

    def manage_resize(self, distro):
        mem, ok = QInputDialog.getText(self, "调整内存", "请输入内存大小（如 4GB 或 2GB）：")
        if ok and mem:
            QMetaObject.invokeMethod(self.backend, "resize_memory", Qt.QueuedConnection, Q_ARG(str, distro), Q_ARG(str, mem))

    def manage_move(self, distro):
        new_dir = QFileDialog.getExistingDirectory(self, f"选择 {distro} 的新位置")
        if new_dir:
            QMetaObject.invokeMethod(self.backend, "move_distro", Qt.QueuedConnection, Q_ARG(str, distro), Q_ARG(str, new_dir))

    def manage_user(self, distro):
        user, ok = QInputDialog.getText(self, "设置默认用户", "请输入用户名（如 root 或普通用户）：")
        if ok and user:
            QMetaObject.invokeMethod(self.backend, "set_default_user", Qt.QueuedConnection, Q_ARG(str, distro), Q_ARG(str, user))

    def run_as_user(self):
        distro = self.current_distro()
        user = self.user_line.text().strip()
        if not distro:
            QMessageBox.warning(self, "提示", "请先选中一个发行版！")
            return
        if not user:
            QMessageBox.warning(self, "提示", "请输入要运行的用户名")
            return
        QMetaObject.invokeMethod(self.backend, "run_as_user", Qt.QueuedConnection, Q_ARG(str, distro), Q_ARG(str, user))

    # -------------- 其余功能 --------------
    def enter(self):
        name = self.current_distro()
        if name:
            subprocess.Popen(["cmd", "/c", "start", "wsl", "-d", name])

    def terminate(self):
        name = self.current_distro()
        if name:
            QMetaObject.invokeMethod(self.backend, "terminate", Qt.QueuedConnection, Q_ARG(str, name))
            QTimer.singleShot(500, self.refresh)

    def backup(self):
        name = self.current_distro()
        if not name:
            return
        default = str(pathlib.Path.home() / f"{name}_{datetime.datetime.now():%Y%m%d_%H%M%S}.tar")
        path, _ = QFileDialog.getSaveFileName(self, "备份到", default, "Tar 文件 (*.tar)")
        if path:
            QMetaObject.invokeMethod(self.backend, "export_distro", Qt.QueuedConnection, Q_ARG(str, name), Q_ARG(str, path))

    def import_distro(self):
        tar, _ = QFileDialog.getOpenFileName(self, "选择备份 tar/tar.gz", "", "Tar 文件 (*.tar *.tar.gz)")
        if not tar:
            return
        name, ok = QInputDialog.getText(self, "导入发行版", "给新发行版起个名字：")
        if not ok or not name:
            return
        install_dir = pathlib.Path.home() / "WSLImport" / name
        install_dir.mkdir(parents=True, exist_ok=True)
        QMetaObject.invokeMethod(self.backend, "import_distro", Qt.QueuedConnection,
                                 Q_ARG(str, name), Q_ARG(str, str(install_dir)), Q_ARG(str, tar))
        QTimer.singleShot(1000, self.refresh)

    def shutdown(self):
        subprocess.run(["wsl", "--shutdown"])
        QTimer.singleShot(500, self.refresh)

    def unregister(self):
        name = self.current_distro()
        if not name:
            return
        if QMessageBox.question(self, "确认", f"即将卸载【{name}】并删除根文件系统，不可恢复！") == QMessageBox.Yes:
            QMetaObject.invokeMethod(self.backend, "unregister", Qt.QueuedConnection, Q_ARG(str, name))
            QTimer.singleShot(500, self.refresh)

    def run_custom_cmd(self):
        cmd = self.cmd_line.text().strip()
        if not cmd:
            QMessageBox.warning(self, "提示", "请输入要执行的命令")
            return
        QMetaObject.invokeMethod(self.backend, "run_cmd", Qt.QueuedConnection, Q_ARG(str, cmd))
        self.cmd_line.clear()


# ---------------- 启动 ----------------
if __name__ == "__main__":
    app = QApplication(sys.argv)
    app.setQuitOnLastWindowClosed(False)  # 关键：关窗口不退出
    w = MainWindow()
    w.show()
    sys.exit(app.exec())
